#include "MrtProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Game/GameState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "Game/Component/SceneModel.h"
#include "LocalShader.h"
#include "LocalScene.h"
#include "Render/VulkanTool.h"

void MrtProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);
	mSceneEntity = static_cast<LocalEntity*>(mGameScene->GetEntityByName("scene"));

	SceneModel& model = mSceneEntity->GetComponent<SceneModel>();
	mObjectSize = model.mDescs.size();
	mPushConstants.resize(mObjectSize);

	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;
			mUboBuffer = CustomBuffer::Builder()
				.SetData(&uboData)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build(this, "ubo");

			mEmptyImage = Picture2DTexture::Builder()
				.SetPath("Block.png")
				.SetAutoMip()
				.Build(this, "emptyImage");

			mEmptySampler = RHISampler::Builder()
				.SetMaxLod(mEmptyImage->GetMipLevels())
				.Build(this, "emptySampler");


			mSceneEntity->mVertexBuffers.resize(mObjectSize, {});
			mSceneEntity->mIndexBuffers.resize(mObjectSize, {});
			mSceneEntity->mAlbedoImages.resize(mObjectSize, {});
			mSceneEntity->mAlbedoSamplers.resize(mObjectSize, {});
			mSceneEntity->mMrtDescriptorMasks.resize(mObjectSize, {});

			SceneModel& model = mSceneEntity->GetComponent<SceneModel>();

			for (uint32 i = 0; i < mObjectSize; ++i)
			{
				std::vector<LocalVertex> tempVertexs;
				tempVertexs.resize(model.mDescs[i].mVertexs.size());
				for (uint32 j = 0; j < model.mDescs[i].mVertexs.size(); ++j)
				{
					tempVertexs[j].position = model.mDescs[i].mVertexs[j].Position;
					tempVertexs[j].normal = model.mDescs[i].mVertexs[j].Normal;
					tempVertexs[j].uv = model.mDescs[i].mVertexs[j].TexC;
				}

				mSceneEntity->mVertexBuffers[i] = VertexBuffer::Builder()
					.SetData(tempVertexs.data())
					.SetSize(tempVertexs.size() * sizeof(LocalVertex))
					.Build();

				mSceneEntity->mIndexBuffers[i] = IndexBuffer::Builder()
					.SetData(model.mDescs[i].mIndices.data())
					.SetSize(model.mDescs[i].mIndices.size() * sizeof(uint32))
					.SetIndexCount(model.mDescs[i].mIndices.size())
					.Build();

				mPushConstants[i].model = mSceneEntity->GetTransform().GetWorldMatrix();
				mPushConstants[i].albedo = model.mDescs[i].mDiffuse;
				mPushConstants[i].hasTexture = model.mDescs[i].mDiffusePath.length() > 0 ? 1 : 0;

				if (mPushConstants[i].hasTexture == 1)
				{
					mSceneEntity->mAlbedoImages[i] = Picture2DTexture::Builder()
						.SetPath(model.mDescs[i].mDiffusePath)
						.SetAutoMip()
						.Build();

					mSceneEntity->mAlbedoSamplers[i] = RHISampler::Builder()
						.SetMaxLod(mSceneEntity->mAlbedoImages[i]->GetMipLevels())
						.Build();
				}
				else
				{
					mSceneEntity->mAlbedoImages[i] = mEmptyImage;
					mSceneEntity->mAlbedoSamplers[i] = mEmptySampler;
				}
			}


			mDescriptor = RHIDescriptor::Builder()
				.PushShader<MrtVertexShader>()
				.PushShader<MrtFragmentShader>()
				.SetDescriptorCount(0, mObjectSize)
				.Build();

			for (int32 i = 0; i < mObjectSize; ++i)
			{
				mDescriptor->SetValue("ubo", mUboBuffer);
				mDescriptor->SetValue("albedoImage", mSceneEntity->mAlbedoImages[i], mSceneEntity->mAlbedoSamplers[i]);
				mSceneEntity->mMrtDescriptorMasks[i] = mDescriptor->AllocateDescriptor();
			}

			mPositionRenderTarget = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "positionRenderTarget");

			mNormalRenderTarget = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "normalRenderTarget");

			mAlbedoRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "albedoRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
				.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "depthRenderTarget");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mPositionRenderTarget)
				.AddColorRenderTarget(mNormalRenderTarget)
				.AddColorRenderTarget(mAlbedoRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mPositionRenderTarget)
				.AddRenderTarget(mNormalRenderTarget)
				.AddRenderTarget(mAlbedoRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void MrtProcess::OnTick(float deltaTime)
{
	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;
			mUboBuffer->WriteMemory(&uboData);

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			for (int32 i = 0; i < mObjectSize; ++i)
			{
				vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mSceneEntity->mMrtDescriptorMasks[i].GetDescriptorSet(), 0, nullptr);

				VkBuffer vertexBuffer = mSceneEntity->mVertexBuffers[i]->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = mSceneEntity->mIndexBuffers[i]->GetBuffer();
				vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConstants), &mPushConstants[i]);

				vkCmdDrawIndexed(commandBuffer, mSceneEntity->mIndexBuffers[i]->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(commandBuffer);
		});
}

void MrtProcess::OnResize()
{
	mPositionRenderTarget = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "positionRenderTarget");

	mNormalRenderTarget = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "normalRenderTarget");

	mAlbedoRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "albedoRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
		.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "depthRenderTarget");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mPositionRenderTarget)
		.AddRenderTarget(mNormalRenderTarget)
		.AddRenderTarget(mAlbedoRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
